Score-based Data Assimilation StormCast by NickGeneva · Pull Request #730 · NVIDIA/earth2studio

NickGeneva · 2026-03-05T21:03:23Z

Earth2Studio Pull Request

Description

Adds stormcast SDA model, a stateful data assimilation method (this is probably one of the more complex implementations of SDA compared to something like corrdiff)
An example demonstrating how to use it with the Integrated Surface Database provided by HRRR
Adds an additional requirement onto the DA model protocol, the concept of init_coords which reflect initial inputs needed for create_generator and the first set of parameters in call
Also update protocol to allow None obs

Coverage:

earth2studio/models/da/__init__.py                2      0   100%
earth2studio/models/da/base.py                    8      0   100%
earth2studio/models/da/interp.py                202    176    13%   86-100, 104, 116, 152-157, 170-174, 202-208, 227-311, 346-380, 397-411, 438-589
earth2studio/models/da/sda_stormcast.py         279     15    95%   285, 386-387, 391-393, 433, 597, 691, 881, 884, 890, 894, 903-906, 909

Rendered Example

_home_ngeneva_Downloads__build__build58_html_examples_21_stormcast_sda html (1)

Closer results

Checklist

I am familiar with the Contributing Guidelines.
New or existing tests cover these changes.
The documentation is up to date with these changes.
The CHANGELOG.md is up to date with these changes.
An issue is linked to this pull request.
Assess and address Greptile feedback (AI code review bot for guidance; use discretion, addressing all feedback is not required).

Dependencies

greptile-apps · 2026-03-05T21:07:19Z

Greptile Summary

This PR introduces StormCastSDA, a stateful score-based data assimilation model that combines a regression network and a diffusion model with DPS (Diffusion Posterior Sampling) guidance to assimilate sparse point observations into convection-allowing regional forecasts. It also extends the AssimilationModel protocol with init_coords() and optional None observation support, adds a legacy=False mode to fetch_data that returns cupy-backed xr.DataArray directly (avoiding unnecessary CPU round-trips), and introduces a GPU-safe map_coords_xr utility.

Key changes:

earth2studio/models/da/sda_stormcast.py: Full SDA StormCast implementation, including KD-tree observation mapping, ray-casting point-in-polygon guard, GPU bilinear conditioning interpolation, and averaged scatter-add for duplicate observations. One open concern: NaN observation values are not filtered before the scatter-add in _build_obs_tensors, which can silently corrupt the DPS guidance if a DataFrame contains missing readings.
earth2studio/data/utils.py: New legacy=False path in fetch_data is well-guarded (device.index or 0); minor: interp_method is silently ignored in this path when interp_to=None.
earth2studio/utils/coords.py: New map_coords_xr correctly handles cupy backends via direct numpy/cupy indexing instead of xarray.interp().
earth2studio/models/da/base.py: Clean protocol extension.
examples/21_stormcast_sda.py: End-to-end example with minor duplicate imports.
Test coverage is thorough, including GPU vs. scipy reference validation for the bilinear interpolation path.

Confidence Score: 4/5

PR is safe to merge with one recommended fix — NaN observation filtering in _build_obs_tensors.
The implementation is well-structured with good test coverage and most previously-raised issues have been addressed. The primary remaining concern is that NaN observation values (common in real weather station data) are not filtered before the scatter-add in _build_obs_tensors, which can silently corrupt the DPS guidance tensor. The interp_method silent-ignore issue is minor. Everything else — GPU path correctness, averaging of duplicate observations, conditioning interpolation, ascending-order guard, docstrings — looks solid.
earth2studio/models/da/sda_stormcast.py — NaN filtering in _build_obs_tensors (line 599).

Important Files Changed

Filename	Overview
earth2studio/models/da/sda_stormcast.py	Core StormCast SDA implementation — large, well-structured new file. NaN observation values are not filtered before scatter-add in `_build_obs_tensors`, which can silently corrupt DPS guidance. Most previously-flagged issues (debug prints, mutable default args, typos, sorting guard, duplicate obs averaging, sampler_args key validation) have been addressed.
earth2studio/data/utils.py	Adds a `legacy` mode to `fetch_data` that returns a raw `xr.DataArray` with cupy backing on CUDA. The `device.index or 0` guard is correctly applied in the new non-legacy path. Minor: `interp_method` is silently ignored in the non-legacy path when `interp_to=None`.
earth2studio/models/da/base.py	Protocol extended to allow `None` observations in `__call__` and `create_generator`, and adds optional `*args` init parameters and `init_coords()` method. Changes are clean and well-documented.
earth2studio/models/da/interp.py	Renames `tolerance` to `time_tolerance` and adds `init_coords()` returning `None`. The rename is a breaking API change but accepted per prior thread discussion. Smolyak interpolation logic is unchanged and looks correct.
earth2studio/utils/coords.py	New `map_coords_xr` function implementing GPU/CPU-aware nearest-neighbor coordinate mapping without calling `xarray.interp()` (avoiding the previous scipy/CPU-only path). Uses sort-based `searchsorted` with correct ascending-order handling via `np.argsort`. Logic looks correct.
examples/21_stormcast_sda.py	Well-written end-to-end example. `cartopy`/`matplotlib` are imported twice (lines 165–167 and 235–237), which is redundant. The `.get()` calls are CUDA-specific but the example is documented as GPU-only. Other previously flagged issues appear to be addressed.
test/models/da/test_da_sda_stormcast.py	Comprehensive test suite covering polygon point-in-polygon, observation tensor building (including None, out-of-grid, and duplicate-averaging cases), conditioning fetch, `__call__`, `create_generator`, and exception handling. GPU interpolation test validates against scipy reference. Good coverage.
test/data/test_data_utils.py	New tests for the `legacy=False` mode of `fetch_data` and updated `prep_data_inputs` behaviour. Coverage looks correct and complete.

_{Last reviewed commit: 45d157f}

earth2studio/data/utils.py

earth2studio/models/da/sda_stormcast.py

NickGeneva · 2026-03-09T16:25:59Z

@greptile-ai

earth2studio/models/da/sda_stormcast.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

.cursor/rules/e2s-009-prognostic-models.mdc

earth2studio/models/da/interp.py

examples/21_stormcast_sda.py

NickGeneva · 2026-03-11T03:09:59Z

/blossom-ci

NickGeneva · 2026-03-11T03:32:57Z

/blossom-ci

NickGeneva · 2026-03-11T05:34:02Z

/blossom-ci

NickGeneva · 2026-03-11T15:05:07Z

/blossom-ci

NickGeneva · 2026-03-11T20:32:39Z

/blossom-ci

NickGeneva added 9 commits March 3, 2026 03:12

Converitng to 2.0 samplers

a6d33a8

Updating stormcast

c26b3b7

SDA running

eb8cedc

Random pytest fixes

458f5b3

adding cupy support in fetch_data

c2e08cf

Xarray bits

65c5ec1

Adding obs

3215a58

generator

69ab063

Updates

95ae396

greptile-apps bot reviewed Mar 5, 2026

View reviewed changes

earth2studio/data/utils.py Show resolved Hide resolved

earth2studio/models/da/sda_stormcast.py Outdated Show resolved Hide resolved

earth2studio/models/da/sda_stormcast.py Outdated Show resolved Hide resolved

earth2studio/models/da/sda_stormcast.py Outdated Show resolved Hide resolved

NickGeneva added 15 commits March 5, 2026 22:10

Fixing iterator

9c57dda

Updates

e3356de

Adding tests

ba0244b

Updating tolerance name

ab1d9d5

Draft example

60bec2b

Updates

c389222

Updates

00ba3f6

Some updates

447ae2f

Example working better

fa28e80

Example update

82093af

Updates

bbd7290

Improvments

467d171

Merge branch 'main' into ngeneva/sda_stormcast

d02baec

Greptile 1

ba75a90

Greptile 2

1f8aa7e

greptile-apps bot reviewed Mar 9, 2026

View reviewed changes

albertocarpentieri reviewed Mar 9, 2026

View reviewed changes

earth2studio/models/da/sda_stormcast.py Outdated Show resolved Hide resolved

Update earth2studio/models/da/sda_stormcast.py

71d1ad0

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

NickGeneva added 7 commits March 10, 2026 06:46

Remove chardet

65d5f01

Testing

7f8a288

remove repeat imports

513e582

revert isd

96f40ae

Fixing interp

720e442

Drop s3fs version

ae57064

Drop s3fs version

f0be113

pzharrington reviewed Mar 10, 2026

View reviewed changes

.cursor/rules/e2s-009-prognostic-models.mdc Outdated Show resolved Hide resolved

pzharrington reviewed Mar 10, 2026

View reviewed changes

earth2studio/models/da/interp.py Show resolved Hide resolved

pzharrington reviewed Mar 10, 2026

View reviewed changes

examples/21_stormcast_sda.py Outdated Show resolved Hide resolved

pzharrington reviewed Mar 10, 2026

View reviewed changes

examples/21_stormcast_sda.py Outdated Show resolved Hide resolved

NickGeneva added 5 commits March 11, 2026 02:18

Merge branch 'main' into ngeneva/sda_stormcast

96ceff9

Feedback

47dbaf4

Feedback

deebc08

Adding das to all install

6958de5

Adding das to all install

7b9415b

NickGeneva added 2 commits March 11, 2026 03:22

Adding install

bdd2760

Lint Fixes

0246b63

NickGeneva added 3 commits March 11, 2026 05:19

Merge branch 'main' into ngeneva/sda_stormcast

dc9e312

fix test

8d2d420

fix test

e0fa6d9

Merge branch 'main' into ngeneva/sda_stormcast

349d9d5

pzharrington approved these changes Mar 11, 2026

View reviewed changes

Merge branch 'main' into ngeneva/sda_stormcast

86e1148

NickGeneva merged commit 53cc0c2 into NVIDIA:main Mar 11, 2026
7 of 10 checks passed

Conversation

NickGeneva commented Mar 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Earth2Studio Pull Request

Description

Checklist

Dependencies

Uh oh!

greptile-apps bot commented Mar 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 4/5

Important Files Changed

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

NickGeneva commented Mar 9, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

NickGeneva commented Mar 11, 2026

Uh oh!

NickGeneva commented Mar 11, 2026

Uh oh!

NickGeneva commented Mar 11, 2026

Uh oh!

NickGeneva commented Mar 11, 2026

Uh oh!

NickGeneva commented Mar 11, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

NickGeneva commented Mar 5, 2026 •

edited

Loading

greptile-apps bot commented Mar 5, 2026 •

edited

Loading